package org.nimi317.web_gis.controller;

import cn.dev33.satoken.annotation.SaCheckRole;
import cn.dev33.satoken.stp.StpUtil;
import cn.hutool.core.util.ObjectUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.metadata.IPage;
import jakarta.validation.Valid;
import org.nimi317.web_gis.Enum.LstmUpdaterEnum;
import org.nimi317.web_gis.Enum.ModelStatus;
import org.nimi317.web_gis.annotation.Page;
import org.nimi317.web_gis.entity.Model;
import org.nimi317.web_gis.entity.R;
import org.nimi317.web_gis.exception.E;
import org.nimi317.web_gis.exception.RException;
import org.nimi317.web_gis.form.post.ModelPagePost;
import org.nimi317.web_gis.form.post.ModelPost;
import org.nimi317.web_gis.form.post.ModelPredictPost;
import org.nimi317.web_gis.form.post.ModelPredictStepPost;
import org.nimi317.web_gis.model.PredictModel;
import org.nimi317.web_gis.service.ModelService;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.util.Arrays;
import java.util.List;
import java.util.Map;

/**
 * 模型(Model)表控制层
 *
 * @author makejava
 * @since 2024-05-13 15:29:23
 */
@RestController
@RequestMapping("model")
public class ModelController {

    /**
     * 服务对象
     */
    private final ModelService modelService;

    public ModelController(ModelService modelService) {
        this.modelService = modelService;
    }

    /**
     * 模型训练
     */
    @PostMapping
    public R<?> train(@RequestPart @Valid ModelPost post, @RequestPart MultipartFile file) {
        modelService.train(post, file);
        return R.ok("模型训练启动成功");
    }

    @PostMapping("predict")
    public R<?> predict(@RequestBody @Valid ModelPredictPost post) {
        Model model = modelService.getById(post.getModelId());
        if (ObjectUtil.isNull(model)) {
            throw new RException(E.ModelNotFoundException);
        }
        if (!model.getStatus().equals(ModelStatus.Finished)) {
            throw new RException(E.ModelNotFinishedException);
        }
        double[][] predict = new PredictModel(model).predict(post.getInput());
        return R.ok(predict);
    }

    @PostMapping("predict/step")
    public R<?> predictStep(@RequestBody @Valid ModelPredictStepPost post) {
        Model model = modelService.getById(post.getModelId());
        if (ObjectUtil.isNull(model)) {
            throw new RException(E.ModelNotFoundException);
        }
        if (!model.getStatus().equals(ModelStatus.Finished)) {
            throw new RException(E.ModelNotFinishedException);
        }
        double[][][] predict = new PredictModel(model).predict(post.getInput(), post.getStep());
        return R.ok(predict);
    }

    @PostMapping("page")
    @Page
    public R<?> train(@RequestBody @Valid ModelPagePost post) {
        IPage<Model> page = modelService.pageModel(post);
        return R.ok(page);
    }

    @GetMapping("dict")
    public R<?> dict() {
        List<? extends Map<?, ?>> update = Arrays.stream(LstmUpdaterEnum.values()).map(LstmUpdaterEnum::getObject).toList();
        List<? extends Map<String, ?>> status = Arrays.stream(ModelStatus.values()).map(ModelStatus::getObject).toList();
        return R.ok(Map.of("update", update, "status", status));
    }

    @GetMapping
    @SaCheckRole("admin")
    public R<?> list() {
        LambdaQueryWrapper<Model> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Model::getStatus, ModelStatus.Finished);
        return R.ok(modelService.list(wrapper));
    }

    @DeleteMapping("{id}")
    public R<?> delete(@PathVariable Long id) {
        LambdaQueryWrapper<Model> wrapper = new LambdaQueryWrapper<>();
        wrapper.eq(Model::getId, id)
                .eq(Model::getUserId, StpUtil.getSession().getLoginId());
        if (modelService.remove(wrapper)) {
            return R.ok("删除成功");
        }
        return R.fail("删除失败模型不存在");
    }
}

